#include "stdafx.h"
#include "A2DImage.h"
#include <setjmp.h>

extern "C" 
{
#	include "jpeglib.h"
#	include "jerror.h"
};

#define BLOCK_SIZE	4096

struct MySourceMgr : jpeg_source_mgr
{
	LPBYTE m_data, m_dataEnd;
	LPBYTE m_dataPtr;
	BYTE m_buffer[BLOCK_SIZE];
	bool m_startOfFile;
};

void jpeg_init_source(j_decompress_ptr cinfo)
{
	MySourceMgr *self = (MySourceMgr*)cinfo->src;
	self->m_dataPtr = self->m_data;
	self->m_startOfFile = true;
}

boolean jpeg_fill_input_buffer(j_decompress_ptr cinfo)
{
	MySourceMgr *self = (MySourceMgr*)cinfo->src;
	
	DWORD bytesLeft = self->m_dataEnd - self->m_dataPtr;	
	DWORD bytesRead = 0;
	if (bytesLeft > 0)
	{
		DWORD bytesToRead = min(bytesLeft, BLOCK_SIZE);
		memcpy_s(self->m_buffer, BLOCK_SIZE, self->m_dataPtr, bytesToRead);
		self->m_dataPtr += bytesToRead;
		bytesRead = bytesToRead;
	}
	else
	{
		if (self->m_startOfFile)
			ERREXIT(cinfo, JERR_INPUT_EMPTY);
		WARNMS(cinfo, JWRN_JPEG_EOF);
		/* Insert a fake EOI marker */
		self->m_buffer[0] = 0xFF;
		self->m_buffer[1] = JPEG_EOI;
		bytesRead = 2;
	}

	self->next_input_byte = self->m_buffer;
	self->bytes_in_buffer = bytesRead;
	self->m_startOfFile = FALSE;
	return TRUE;
}

void jpeg_skip_input_buffer(j_decompress_ptr cinfo, long num_bytes)
{
	MySourceMgr *self = (MySourceMgr*)cinfo->src;

	if (num_bytes > 0)
	{
		while (num_bytes > (long)self->bytes_in_buffer)
		{
			num_bytes -= (long)self->bytes_in_buffer;
			jpeg_fill_input_buffer(cinfo);
		}
		self->next_input_byte += (size_t)num_bytes;
		self->bytes_in_buffer -= (size_t)num_bytes;
	}
}

void jpeg_term_source(j_decompress_ptr cinfo) 
{ 
	// do nothing
}

struct MyErrorMgr : jpeg_error_mgr
{
	jmp_buf m_env;
};

void jpeg_error_exit(j_common_ptr cinfo) 
{
	MyErrorMgr *err = (MyErrorMgr*)cinfo->err;
	longjmp(err->m_env, -1);
}

bool A2DImage::LoadAsJPEG(LPVOID data, DWORD size)
{
	jpeg_decompress_struct cinfo;
	jpeg_create_decompress(&cinfo);

	MyErrorMgr err;
	jpeg_std_error(&err);
	err.error_exit = jpeg_error_exit;
	cinfo.err = &err;

	MySourceMgr src;
	src.init_source = jpeg_init_source;
	src.fill_input_buffer = jpeg_fill_input_buffer;
	src.skip_input_data = jpeg_skip_input_buffer;
	src.resync_to_restart = jpeg_resync_to_restart;
	src.term_source = jpeg_term_source;
	src.m_data = (LPBYTE)data;
	src.m_dataEnd = (LPBYTE)data + size;
	src.next_input_byte = NULL;
	src.bytes_in_buffer = 0;
	cinfo.src = &src;

	if (setjmp(err.m_env) != 0)
	{
		jpeg_destroy_decompress(&cinfo);
		return false;
	}

	if (jpeg_read_header(&cinfo, TRUE) != JPEG_HEADER_OK)
		longjmp(err.m_env, -1);

	if (cinfo.out_color_space == JCS_CMYK || cinfo.out_color_space == JCS_YCCK)
		cinfo.out_color_space = JCS_CMYK;
	else
		cinfo.out_color_space = JCS_RGB;

	if (!jpeg_start_decompress(&cinfo))
		longjmp(err.m_env, -1);

	if (!CreateTexture(cinfo.output_width, cinfo.output_height))
		longjmp(err.m_env, -1);

	LPBYTE dstBuf;
	int dstPitch;
	if (!Lock(dstBuf, dstPitch))
		longjmp(err.m_env, -1);
	LPBYTE dstRow = dstBuf;

	int srcPitch = cinfo.output_width * cinfo.output_components;
	CAutoVectorPtr<BYTE> srcBuf;
	srcBuf.Allocate(srcPitch);
	LPBYTE srcRow = srcBuf;

	if (cinfo.out_color_space == JCS_CMYK)
	{
		assert(cinfo.output_components == 4);

		JSAMPROW jrow[1] = { srcRow };
		for (; cinfo.output_scanline < cinfo.output_height; dstRow += dstPitch)
		{
			jpeg_read_scanlines(&cinfo, jrow, 1);

			D3DCOLOR *dstPixel = (D3DCOLOR*)dstRow;
			BYTE *srcPixel = srcRow;
			for (UINT x = 0; x < cinfo.output_width; x++, dstPixel++, srcPixel += 4)
			{
				*dstPixel = D3DCOLOR_XRGB(
					(srcPixel[3] * srcPixel[0]) / 255, // R
					(srcPixel[3] * srcPixel[1]) / 255, // G
					(srcPixel[3] * srcPixel[2]) / 255  // B
				);
			}
		}
	}
	else 
	{
		assert(cinfo.output_components == 3);

		JSAMPROW jrow[1] = { srcRow };
		for (; cinfo.output_scanline < cinfo.output_height; dstRow += dstPitch)
		{
			jpeg_read_scanlines(&cinfo, jrow, 1);

			D3DCOLOR *dstPixel = (D3DCOLOR*)dstRow;
			BYTE *srcPixel = srcRow;
			for (size_t x = 0; x < cinfo.output_width; x++, dstPixel++, srcPixel += 3)
			{
				*dstPixel = D3DCOLOR_XRGB(
					srcPixel[0], // R
					srcPixel[1], // G
					srcPixel[2]  // B
				);
			}
		}
	}

	Unlock();
	jpeg_finish_decompress(&cinfo);
	jpeg_destroy_decompress(&cinfo);
	return true;
}